import tensorflow as tf
import os
import sys
from python_ai.common.xcommon import *
import numpy as np

model_path = '../../../../../../large_data/model/inceptionV3/tensorflow_inception_graph.pb'
IMG_DIR = '../../../../../../large_data/CV3/_many_files/agriculture/train'
VER = 'v2.0'
FILE_NAME = os.path.basename(__file__)
OUTPUT_DIR = os.path.join('_save', FILE_NAME, VER)
VEC_SAVE_DIR = os.path.join(OUTPUT_DIR, 'vectors')
os.makedirs(VEC_SAVE_DIR, exist_ok=True)

with open(model_path, 'rb') as f:
    bin = f.read()

graphDef = tf.GraphDef()
graphDef.ParseFromString(bin)

# DecodeJpeg: TypeError: Cannot interpret feed_dict key as Tensor: Can not convert a Operation into a Tensor.
# DecodeJpeg/contents: TypeError: Cannot interpret feed_dict key as Tensor: Can not convert a Operation into a Tensor.
# DecodeJpeg/contents:0 OK
# pool_3/_reshape Get None
# pool_3/_reshape:0 OK
input_placeholder, output_tensor = tf.import_graph_def(graphDef,
                                                       return_elements=[
                                                           'DecodeJpeg/contents:0',
                                                           'pool_3/_reshape:0'
                                                       ])

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    print('Processing img ...')
    cnt = 0
    for sub_dir in os.listdir(IMG_DIR):
        sub_dir_path = os.path.join(IMG_DIR, sub_dir)
        vec_dir_path = os.path.join(VEC_SAVE_DIR, sub_dir)
        os.makedirs(vec_dir_path, exist_ok=True)
        for file in os.listdir(sub_dir_path):
            cnt += 1
            if cnt % 25 == 0:
                print(f'Processing img no.{cnt} ...')
            vec_path = os.path.join(vec_dir_path, file + '.txt')
            if os.path.exists(vec_path):
                continue
            file_path = os.path.join(sub_dir_path, file)
            with open(file_path, 'rb') as f:
                img_bin = f.read()
            vec = sess.run(output_tensor, feed_dict={input_placeholder: img_bin})
            np.savetxt(vec_path, vec)
    print('All img processed')
